#!/usr/bin/env python3
"""
服务停止脚本 - 查找并终止uvicorn进程
"""
import os
import sys
import signal
import subprocess
import time
import psutil

def find_uvicorn_processes():
    """查找所有uvicorn进程"""
    uvicorn_processes = []
    
    for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
        try:
            # 检查进程名和命令行参数
            if proc.info['name'] == 'python' or proc.info['name'] == 'python.exe':
                cmdline = proc.info['cmdline']
                if cmdline and ('uvicorn' in ' '.join(cmdline)):
                    uvicorn_processes.append(proc)
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass
    
    return uvicorn_processes

def graceful_shutdown(pid):
    """尝试优雅地关闭进程"""
    try:
        process = psutil.Process(pid)
        print(f"尝试优雅地关闭进程 {pid}...")
        
        # 首先尝试发送SIGTERM信号
        process.send_signal(signal.SIGTERM)
        
        # 等待进程结束，最多10秒
        for _ in range(10):
            if not process.is_running():
                print(f"进程 {pid} 已停止")
                return True
            time.sleep(1)
        
        # 如果进程仍在运行，强制关闭
        print(f"进程 {pid} 未响应，强制终止...")
        process.kill()
        return True
    
    except psutil.NoSuchProcess:
        print(f"进程 {pid} 不存在")
        return True
    except Exception as e:
        print(f"关闭进程 {pid} 失败: {e}")
        return False

def force_kill(pid):
    """强制终止进程"""
    try:
        process = psutil.Process(pid)
        print(f"强制终止进程 {pid}...")
        process.kill()
        return True
    except Exception as e:
        print(f"终止进程 {pid} 失败: {e}")
        return False

def main():
    """主函数"""
    # 检查是否安装了psutil
    try:
        import psutil
    except ImportError:
        print("错误: 未安装psutil模块。请先安装:")
        print("pip install psutil")
        return False
    
    print("查找uvicorn进程...")
    uvicorn_processes = find_uvicorn_processes()
    
    if not uvicorn_processes:
        print("未找到运行中的uvicorn进程")
        return False
    
    print(f"找到 {len(uvicorn_processes)} 个uvicorn进程")
    
    all_stopped = True
    for proc in uvicorn_processes:
        pid = proc.info['pid']
        cmdline = ' '.join(proc.info['cmdline'])
        print(f"进程 {pid}: {cmdline}")
        
        # 尝试优雅关闭
        if not graceful_shutdown(pid):
            # 如果失败，尝试强制终止
            if not force_kill(pid):
                all_stopped = False
                print(f"无法终止进程 {pid}")
    
    if all_stopped:
        print("所有uvicorn进程已停止")
    else:
        print("部分进程无法停止，请尝试手动终止")
    
    return all_stopped

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)